[red-knot] Optimise visibility constraints for `*`-import definitions #17317

AlexWaygood · 2025-04-09T15:30:57Z

Summary

@sharkdp gets all the credit for this idea! Quoting his comments from Discord:

I have an idea for an optimization: It is kind of wasteful to generate the full
if <placeholder for sym1>:
    from module import sym1
else:
    # implicit empty else branch

if <placeholder for sym2>:
    from module import sym2
else:
    # implicit empty else branch


# …
for all N symbols that we would *-import from module. One crucial difference w.r.t. normal control flow is that we already know which definitions will appear inside the branches. So I think we could check if we already have a definition for sym_i prior to the star-import. If we do, we proceed as before. But if we do not (which should be the overwhelming majority of cases), we can potentially generate something simpler. It's probably enough to only record that from module import sym_i definition and apply the placeholder visibility constraint to just that definition (instead of all active definitions). This might allow us to get rid of snapshotting completely for the majority of cases.

Test Plan

Existing tests all pass, and Codspeed is very happy about the PR!

github-actions · 2025-04-09T15:32:48Z

`mypy_primer` results

No ecosystem changes detected ✅

codspeed-hq · 2025-04-09T15:38:04Z

CodSpeed Performance Report

Merging #17317 will improve performances by 10.3%

_{Comparing alex/star-import-optimisation (0bc6440) with main (ff376fc)}

Summary

⚡ 1 improvements
✅ 31 untouched benchmarks

Benchmarks breakdown

	Benchmark	`BASE`	`HEAD`	Change
⚡	`red_knot_check_file[cold]`	111.4 ms	101 ms	+10.3%

carljm

Nice!

carljm · 2025-04-09T15:59:15Z

crates/red_knot_python_semantic/src/semantic_index/builder.rs

+                                let constraint_id = self
+                                    .current_use_def_map_mut()
+                                    .record_star_import_visibility_constraint(
+                                        star_import,
+                                        symbol_id,
+                                    );
+
+                                let post_definition = self.flow_snapshot();
+                                self.flow_restore(pre_definition);
+
+                                self.current_use_def_map_mut()
+                                    .negate_star_import_visibility_constraint(
+                                        symbol_id,
+                                        constraint_id,
+                                    );
+
+                                self.flow_merge(post_definition);


This PR reclaims almost all the lost perf, so it seems fine to just go ahead with it as-is. But it still seems like we do more work here than necessary? In the newly-added case, we shouldn't have to do any snapshotting or flow-state merging at all. We can just add the star-import definition, apply the (positive) visibility constraint to it, and move on.

As in, something like this (relative to my PR)? Quite a lot of tests start failing if I apply this diff...

diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index 486482f7f..406014831 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -1219,23 +1219,12 @@ where // we can apply the visibility constraint to *only* the added definition, // rather than all definitions if newly_added { - let constraint_id = self + self .current_use_def_map_mut() .record_star_import_visibility_constraint( star_import, symbol_id, ); - - let post_definition = self.flow_snapshot(); - self.flow_restore(pre_definition); - - self.current_use_def_map_mut() - .negate_star_import_visibility_constraint( - symbol_id, - constraint_id, - ); - - self.flow_merge(post_definition); } else { let constraint_id = self.record_visibility_constraint(star_import.into()); diff --git a/crates/red_knot_python_semantic/src/semantic_index/use_def.rs b/crates/red_knot_python_semantic/src/semantic_index/use_def.rs index edb27d26e..002faf319 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/use_def.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/use_def.rs @@ -713,19 +713,6 @@ impl<'db> UseDefMapBuilder<'db> { visibility_id } - pub(super) fn negate_star_import_visibility_constraint( - &mut self, - symbol_id: ScopedSymbolId, - constraint: ScopedVisibilityConstraintId, - ) { - let negated_constraint = self.visibility_constraints.add_not_constraint(constraint); - self.symbol_states[symbol_id] - .record_visibility_constraint(&mut self.visibility_constraints, negated_constraint); - self.scope_start_visibility = self - .visibility_constraints - .add_and_constraint(self.scope_start_visibility, negated_constraint); - } - /// This method resets the visibility constraints for all symbols to a previous state /// *if* there have been no new declarations or bindings since then. Consider the /// following example:

I'll go ahead and merge for now, but feel free to file a followup PR if you can see a way of optimising this further without breaking all my tests ;)

This is something I also didn't think about when writing that message on Discord, but I think you would also need to apply the negated visibility constraint to the implicit symbol = <unbound> binding (it is stored as a None entry at the beginning of the vector). I think that should allow you to get rid of snapshotting entirely for the "fast path".

We would basically still model

if <placeholder>: from module import symbol else: symbol = <unbound>

but with the advantage that the constraint does not need to be applied to anything else.

Ah yeah, that's because "unbound" is a binding also, and we do need to apply the negated visibility constraint to the "unbound" binding. (The "unbound" binding is visible only if the star import target symbol does not end up existing.) I think this should still be doable by just applying constraints to the right bindings, without snapshotting and merging, but it would require even more special-cased new APIs. Not sure it's worth it since this version already does so well, definitely no need to look into it now.

Okay -- I'm going to move on for now as I think both of you have a far better understanding of our visibility constraints machinery than I do. But I welcome further optimisations of this code from anybody who wants to work on them 😄

Not sure it's worth it since this version already does so well, definitely no need to look into it now.

I agree. Seems like the new isolated application of constraints helps enough to simplify the overall constraint structure.

…astral-sh#17317)

[red-knot] Optimise visibility constraints for *-import definitions

5caf5b4

AlexWaygood added the ty Multi-file analysis & type inference label Apr 9, 2025

AlexWaygood marked this pull request as ready for review April 9, 2025 15:38

AlexWaygood requested review from carljm, sharkdp and dcreager as code owners April 9, 2025 15:38

AlexWaygood mentioned this pull request Apr 7, 2025

[red-knot] support * imports #14169

Closed

AlexWaygood added the performance Potential performance improvement label Apr 9, 2025

carljm approved these changes Apr 9, 2025

View reviewed changes

improve type safety

0bc6440

AlexWaygood force-pushed the alex/star-import-optimisation branch from 5730000 to 0bc6440 Compare April 9, 2025 16:50

AlexWaygood enabled auto-merge (squash) April 9, 2025 16:50

AlexWaygood merged commit 7339902 into main Apr 9, 2025
20 checks passed

AlexWaygood deleted the alex/star-import-optimisation branch April 9, 2025 16:53

Glyphack pushed a commit to Glyphack/ruff that referenced this pull request Apr 9, 2025

[red-knot] Optimise visibility constraints for *-import definitions (…

34a436c

…astral-sh#17317)

AlexWaygood mentioned this pull request Apr 9, 2025

[red-knot] Consider further optimising visibility constraints applied to *-import definitions #17322

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[red-knot] Optimise visibility constraints for `*`-import definitions #17317

[red-knot] Optimise visibility constraints for `*`-import definitions #17317

Uh oh!

AlexWaygood commented Apr 9, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Apr 9, 2025 •

edited

Loading

Uh oh!

codspeed-hq bot commented Apr 9, 2025 •

edited

Loading

Uh oh!

carljm left a comment

Uh oh!

carljm Apr 9, 2025 •

edited

Loading

Uh oh!

AlexWaygood Apr 9, 2025

Uh oh!

AlexWaygood Apr 9, 2025

Uh oh!

sharkdp Apr 9, 2025

Uh oh!

carljm Apr 9, 2025

Uh oh!

AlexWaygood Apr 9, 2025

Uh oh!

sharkdp Apr 9, 2025

Uh oh!

Uh oh!

Uh oh!

[red-knot] Optimise visibility constraints for *-import definitions #17317

[red-knot] Optimise visibility constraints for *-import definitions #17317

Uh oh!

Conversation

AlexWaygood commented Apr 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test Plan

Uh oh!

github-actions bot commented Apr 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

mypy_primer results

Uh oh!

codspeed-hq bot commented Apr 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

CodSpeed Performance Report

Merging #17317 will improve performances by 10.3%

Summary

Benchmarks breakdown

Uh oh!

carljm left a comment

Choose a reason for hiding this comment

Uh oh!

carljm Apr 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AlexWaygood Apr 9, 2025

Choose a reason for hiding this comment

Uh oh!

AlexWaygood Apr 9, 2025

Choose a reason for hiding this comment

Uh oh!

sharkdp Apr 9, 2025

Choose a reason for hiding this comment

Uh oh!

carljm Apr 9, 2025

Choose a reason for hiding this comment

Uh oh!

AlexWaygood Apr 9, 2025

Choose a reason for hiding this comment

Uh oh!

sharkdp Apr 9, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

[red-knot] Optimise visibility constraints for `*`-import definitions #17317

[red-knot] Optimise visibility constraints for `*`-import definitions #17317

AlexWaygood commented Apr 9, 2025 •

edited

Loading

github-actions bot commented Apr 9, 2025 •

edited

Loading

`mypy_primer` results

codspeed-hq bot commented Apr 9, 2025 •

edited

Loading

carljm Apr 9, 2025 •

edited

Loading